#include "loc_correction/optimizecampos.h"


double computeLoss(const Eigen::Vector3d& cam_pos, 
                   const Eigen::Quaterniond& cam_orientation,
                   const Eigen::Matrix4d& pnp_mat,
                   const Eigen::Matrix4d& ref_uav_pose) {
    // 构造相机变换矩阵 cam_mat
    Eigen::Matrix4d cam_mat = Eigen::Matrix4d::Identity();
    cam_mat.block<3, 1>(0, 3) = cam_pos;
    cam_mat.block<3, 3>(0, 0) = cam_orientation.toRotationMatrix();

    // 计算估计的无人机位姿
    Eigen::Matrix4d est_uav_pose = pnp_mat * cam_mat;

    // 计算位置误差
    Eigen::Vector3d position_error = est_uav_pose.block<3, 1>(0, 3) - ref_uav_pose.block<3, 1>(0, 3);//融合值
    double position_loss = position_error.squaredNorm();

    // 计算姿态误差
    Eigen::Quaterniond q_est(est_uav_pose.block<3, 3>(0, 0));
    Eigen::Quaterniond q_ref(ref_uav_pose.block<3, 3>(0, 0));
    double orientation_loss = q_est.angularDistance(q_ref);

    // 返回总误差（可以根据需要对位置误差和姿态误差加权）
    return position_loss + orientation_loss;
}


void optimizeCamPose(Eigen::Vector3d& cam_pos, 
                     Eigen::Quaterniond& cam_orientation,
                     const Eigen::Matrix4d& pnp_mat,
                     const Eigen::Matrix4d& ref_uav_pose,
                     int max_iterations, 
                     double learning_rate) {
    for (int i = 0; i < max_iterations; ++i) {
        // 计算当前误差
        double current_loss = computeLoss(cam_pos, cam_orientation, pnp_mat, ref_uav_pose);
        std::cout << "Iteration " << i << ", Loss: " << current_loss << std::endl;

        // 数值计算位置的梯度
        Eigen::Vector3d grad_position;
        double epsilon = 1e-6;
        for (int j = 0; j < 3; ++j) {
            Eigen::Vector3d pos_offset = cam_pos;
            pos_offset[j] += epsilon;
            double loss_offset = computeLoss(pos_offset, cam_orientation, pnp_mat, ref_uav_pose);
            grad_position[j] = (loss_offset - current_loss) / epsilon;
        }

        // 数值计算旋转的梯度
        Eigen::Vector4d grad_orientation;
        for (int j = 0; j < 4; ++j) {
            Eigen::Quaterniond q_offset = cam_orientation;
            Eigen::Vector4d q_vec = q_offset.coeffs();
            q_vec[j] += epsilon;
            q_offset.coeffs() = q_vec;
            q_offset.normalize();

            double loss_offset = computeLoss(cam_pos, q_offset, pnp_mat, ref_uav_pose);
            grad_orientation[j] = (loss_offset - current_loss) / epsilon;
        }

        // 更新 cam_pos 和 cam_orientation
        cam_pos -= learning_rate * grad_position;
        Eigen::Vector4d q_vec = cam_orientation.coeffs();
        q_vec -= learning_rate * grad_orientation;
        cam_orientation.coeffs() = q_vec;
        cam_orientation.normalize(); 

        
        if (grad_position.norm() < 1e-6 && grad_orientation.norm() < 1e-6) {
            std::cout << "Converged at iteration " << i << std::endl;
            break;
        }
    }
}